/*
 *  Copyright 2013 Utkin Dmitry
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

#include <QDebug>
#include <QHash>
#include <QWaitCondition>
#include <QMutex>
#include <google/protobuf/descriptor.h>
#include <google/protobuf/descriptor.pb.h>
#include "RpcServer.h"
#include "RpcServerSession.h"
#include "protoqtrpc/rpcmessage.pb.h"
#include "protoqtrpc/options.pb.h"
#include "RpcController.h"

namespace protoqtrpc
{

struct RpcServerSession::ActiveRequest
{
    quint64 id;
    ::google::protobuf::Message* request;
    ::google::protobuf::Message* response;
    RpcController controller;
    bool isMultiResponse;

    ActiveRequest():
        id(0), request(NULL), response(NULL), isMultiResponse(false)
    {
    }

    ~ActiveRequest()
    {
        delete request;
        delete response;
    }
};

class RpcServerSessionPrivate
{
public:
    RpcServer* server;
    QHash<quint64, RpcServerSession::ActiveRequest*> requests;
    QMutex mutex;

    RpcServerSessionPrivate(RpcServer* svr):
        server(svr),
        mutex(QMutex::Recursive)
    {
    }
};


RpcServerSession::RpcServerSession(RpcServer* server):
    d_ptr(new RpcServerSessionPrivate(server))
{
    qRegisterMetaType<protoqtrpc::PeerInfo>("protoqtrpc::PeerInfo");
}

RpcServerSession::~RpcServerSession()
{
    if (!d_ptr->requests.isEmpty())
    {
        d_ptr->mutex.lock();
        foreach (RpcServerSession::ActiveRequest* request, d_ptr->requests)
            request->controller.StartCancel();
        d_ptr->mutex.unlock();

        QMutex mutex;
        QWaitCondition wait;
        mutex.lock();
        for (int i = 0; i < 30; ++i)
        {
            wait.wait(&mutex, 100);
            if (d_ptr->requests.isEmpty())
                break;
        }

        if (!d_ptr->requests.isEmpty())
            qWarning() << "Some requests are still running";
    }

    delete d_ptr;
}

void RpcServerSession::handleMessage(const RpcMessage& msg)
{
    switch (msg.type())
    {
    case RpcMessage::RPC_REQUEST:
        handleRequest(msg);
        break;

    case RpcMessage::RPC_CANCEL:
        handleCancel(msg);
        break;

    default:
        qWarning() << "RPC type (" << RpcMessage::RpcType_Name(msg.type()).c_str()
                   << ") handling is not supported";
    }
}

void RpcServerSession::handleRequest(const RpcMessage& msg)
{
    Q_ASSERT(d_ptr->server);

    ActiveRequest* req = new ActiveRequest;
    req->id = msg.id();

    if (!msg.has_service_name() || !msg.has_method_name())
    {
        qWarning() << "Service name or Method name is not set";
        req->controller.SetFailed("Service name or Method name is not set");
        sendResponse(req);
        return;
    }


    ::google::protobuf::Service* service =
            d_ptr->server->getService(msg.service_name());
    if (!service)
    {
        qWarning() << "Couldn't find RPC service for "
                   << msg.service_name().c_str() << ":"
                   << msg.method_name().c_str();

        req->controller.SetFailed("Service \"" + msg.service_name() + "\" not found");
        sendResponse(req);
        return;
    }


    const ::google::protobuf::MethodDescriptor* method =
            service->GetDescriptor()->FindMethodByName(msg.method_name());

    if (!method)
    {
        qWarning() << "Received RPC for method "
                   << msg.service_name().c_str() << ":" << msg.method_name().c_str()
                   << " which has registered a service type but doesn't exist!";

        req->controller.SetFailed("Method \"" + msg.method_name() + "\" in service \"" +
                                  msg.service_name() + "\" not found");
        sendResponse(req);
        return;
    }

    req->response = service->GetResponsePrototype(method).New();
    req->request = service->GetRequestPrototype(method).New();
    req->controller.setPeerInfo(&peerInfo());

    if (msg.has_payload())
    {
        if (!req->request->ParseFromString(msg.payload()))
        {
            std::string error;
            std::vector<std::string> fields;
            req->request->FindInitializationErrors(&fields);
            foreach (const std::string& field, fields)
            {
                if (!error.empty())
                    error += ", ";
                error += field;
            }
            if (!error.empty())
                error.insert(0, "because it is missing required fields: ");

            req->controller.SetFailed("Can't parse message of type \"" +
                                      req->request->GetTypeName() + "\" " + error + ".");
            sendResponse(req);
            return;
        }
    }

    req->isMultiResponse = method->options().GetExtension(protoqtrpc::multiresponse);

    d_ptr->mutex.lock();
    d_ptr->requests[req->id] = req;
    d_ptr->mutex.unlock();


    ::google::protobuf::Closure* closure = (req->isMultiResponse)
            ? ::google::protobuf::NewPermanentCallback(this, &RpcServerSession::sendResponse, req)
            : ::google::protobuf::NewCallback(this, &RpcServerSession::sendResponse, req);

    service->CallMethod(method, &req->controller, req->request, req->response,
                        closure);
}

void RpcServerSession::handleCancel(const RpcMessage& msg)
{
    QMutexLocker l(&d_ptr->mutex);

    ActiveRequest* req = d_ptr->requests.value(msg.id());
    if (req)
        req->controller.StartCancel();
}

void RpcServerSession::sendResponse(ActiveRequest* req)
{
    bool remove = (!req->isMultiResponse || req->controller.IsCanceled());
    if (remove)
    {
        d_ptr->mutex.lock();
        d_ptr->requests.remove(req->id);
        d_ptr->mutex.unlock();
    }

    RpcMessage msg;

    msg.set_id(req->id);
    msg.set_type(RpcMessage::RPC_RESPONSE);

    if (req->controller.IsCanceled())
    {
        msg.set_payload("CANCELLED");
        msg.set_status(RpcMessage::STATUS_CANCELLED);
    }
    else if (req->controller.Failed())
    {
        msg.set_payload(req->controller.ErrorText());
        msg.set_status(RpcMessage::STATUS_ERROR);
    }
    else
    {
        Q_ASSERT(req->response);
        msg.set_payload(req->response->SerializeAsString());
        msg.set_status(RpcMessage::STATUS_SUCCESS);
    }

    emit sendMessage(msg);

    if (remove)
        delete req;
}

}
